#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/9/11 13:42
# @Author  : 王凯
# @File    : es_company_tools.py
# @Project : scrapy_spider
import time

import requests
import retrying
from utils.tools import log


class EsCompanyTools:
    def __init__(self):
        self.url = "https://api-dev.local.wfq2020.com/third-party-api/es/single"
        # self.url = "http://127.0.0.1:8000/es/single"

    @retrying.retry(stop_max_attempt_number=3)
    def _get_company_name_or_taxpayer_id_from_es(self, company_name_list, table='dim_enterprise_mapping_v1_1_0_full_1d'):
        cond = [
                   [
                       "term",
                       "company_name",
                       company_name.replace("（", "(").replace("）", ")")
                   ] for company_name in company_name_list
               ] + [
                   [
                       "term",
                       "company_used_name",
                       company_name.replace("（", "(").replace("）", ")")
                   ] for company_name in company_name_list
               ]
        params = {
            "index": table,
            "database": "dw",
            "filters": [
                [
                    "term",
                    "valid",
                    "1"
                ]
            ],
            "fields": "taxpayer_id,company_used_name,company_name,taxpayer_industry1,taxpayer_industry2,taxpayer_industry3,taxpayer_industry4",
            "shoulds": cond,
            "channel_id": 4087,
            "limit": 10000
        }
        start = time.time() * 1000
        response = requests.post(self.url, json=params, timeout=90)
        msg = f"table:{table},cost time {int(time.time() * 1000 - start)}ms,response:{response.text[:200]}"
        log.info(msg)
        response = response.json()
        datas = response.get("data", {}).get("data", [])

        taxpayer_id_mapping = {}
        company_name_mapping = {}
        industry_mapping = {}
        used_name_mapping = {}
        for data in datas:
            company_used_name = data.get("company_used_name")
            company_name = data.get("company_name")
            industry = data.get("taxpayer_industry4") or data.get("taxpayer_industry3") or data.get("taxpayer_industry2") or data.get("taxpayer_industry1")

            if company_used_name:
                for i in company_used_name:
                    if i in company_name_list:
                        taxpayer_id_mapping[i] = data.get("taxpayer_id")
                        company_name_mapping[i] = company_name
                        used_name_mapping[i] = ','.join(data.get("company_used_name", []))
                        industry_mapping[i] = industry
                    else:
                        company_src = i.replace("(", "（").replace(")", "）")
                        if company_src in company_name_list:
                            taxpayer_id_mapping[company_src] = data.get("taxpayer_id")
                            company_name_mapping[company_src] = company_name
                            used_name_mapping[company_src] = ','.join(data.get("company_used_name") or [])
                            industry_mapping[company_src] = industry

            if company_name in company_name_list:
                taxpayer_id_mapping[company_name] = data.get("taxpayer_id")
                company_name_mapping[company_name] = company_name
                industry_mapping[company_name] = industry
                used_name_mapping[company_name] = ','.join(data.get("company_used_name") or [])
            else:
                company_src = company_name.replace("(", "（", ).replace(")", "）", )
                if company_src in company_name_list:
                    taxpayer_id_mapping[company_src] = data.get("taxpayer_id")
                    company_name_mapping[company_src] = company_name
                    industry_mapping[company_src] = industry
                    used_name_mapping[company_src] = ','.join(data.get("company_used_name") or [])

        return taxpayer_id_mapping, company_name_mapping, industry_mapping, used_name_mapping

    def get_company_name_or_taxpayer_id(self, company_name_list):
        taxpayer_id_mapping, company_name_mapping, industry_mapping, used_name_mapping = self._get_company_name_or_taxpayer_id_from_es(company_name_list, table='dim_enterprise_mapping_v1_1_0_full_1d')
        taxpayer_id_mapping_ind, company_name_mapping_ind, industry_mapping_ind, used_name_mapping_ind = self._get_company_name_or_taxpayer_id_from_es(company_name_list,
                                                                                                                                                       table='dim_individual_business_mapping_v1_1_0_full_1d')
        taxpayer_id_mapping.update(taxpayer_id_mapping_ind)
        company_name_mapping.update(company_name_mapping_ind)
        industry_mapping.update(industry_mapping_ind)
        used_name_mapping.update(used_name_mapping_ind)
        return {
            "company_name_mapping": company_name_mapping,
            "taxpayer_id_mapping": taxpayer_id_mapping,
            "industry_mapping": industry_mapping,
            "used_name_mapping": used_name_mapping,
        }


if __name__ == '__main__':
    es_company_tools = EsCompanyTools()
    print(es_company_tools.get_company_name_or_taxpayer_id([
        '三亚森奥电梯有限公司', '三亚天立方商贸有限公司', '和政县腾胜家政服务有限公司', '和政县旭红动物保健有限责任公司', '浏阳市城东烟花辅助材料加工厂', '长沙市首雅园林有限公司',
        '湖南美奥钾业有限责任公司', '长沙博鑫出口烟花有限公司', '长沙三和管桩有限公司', '长沙雅诚工艺品有限公司', '长沙欧化科技有限公司', '长沙博丰机械有限公司', '长沙海川自动化设备有限公司',
        '长沙科华东菱医疗器械有限公司', '长沙宝莱机械设备进出口有限公司', '湖南胜大钢铁贸易有限公司', '长沙市维旭文化传播有限公司', '中仕建工集团有限公司', '怀化高通交电有限公司',
        '湖南金羽航科技有限公司', '湘潭市钰钢贸易有限公司'
    ]))
